"""
Observation Operators Demo

Demonstrates the use of comprehensive observation operators for GSI analysis.

This demo shows:
1. Creating synthetic atmospheric fields
2. Setting up observation locations
3. Computing forward operators for various observation types
4. Computing innovation vectors
5. Computing adjoint operators for gradient computation
"""

println("="^70)
println("OBSERVATION OPERATORS DEMO")
println("="^70)
println()

# Setup paths
demo_dir = @__DIR__
package_root = dirname(demo_dir)
src_dir = joinpath(package_root, "src")

# Load modules
include(joinpath(src_dir, "ObservationOperators/Interpolation.jl"))
include(joinpath(src_dir, "ObservationOperators/ConventionalObs.jl"))

using .Interpolation
using .ConventionalObs
using LinearAlgebra
using Printf

# Import specific functions
import .ConventionalObs: temperature_forward, temperature_adjoint!
import .ConventionalObs: wind_forward, wind_adjoint!
import .ConventionalObs: humidity_forward, humidity_adjoint!
import .ConventionalObs: surface_pressure_forward, surface_pressure_adjoint!

println("✓ Loaded observation operator modules")
println()

# ============================================================================
# Create synthetic atmospheric state
# ============================================================================

println("Step 1: Creating synthetic atmospheric state")
println("-" * "="^68)

# Grid definition
nx, ny, nz = 20, 15, 10
lon_grid = collect(range(-100.0, -80.0, length=nx))
lat_grid = collect(range(30.0, 45.0, length=ny))
sigma_levels = collect(range(0.1, 1.0, length=nz))

println("  Grid size: $(nx) x $(ny) x $(nz)")
println("  Longitude range: $(lon_grid[1]) to $(lon_grid[end])")
println("  Latitude range: $(lat_grid[1]) to $(lat_grid[end])")
println()

# Create synthetic atmospheric fields
t_field = zeros(nx, ny, nz)
u_field = zeros(nx, ny, nz)
v_field = zeros(nx, ny, nz)
q_field = zeros(nx, ny, nz)

for i in 1:nx, j in 1:ny, k in 1:nz
    # Temperature decreases with height
    t_field[i,j,k] = 290.0 - 50.0 * (1.0 - sigma_levels[k]) +
                     5.0 * sin(2π * lon_grid[i]/20)

    # Westerly wind increasing with height
    u_field[i,j,k] = 10.0 + 20.0 * (1.0 - sigma_levels[k])

    # Meridional component
    v_field[i,j,k] = 5.0 * cos(2π * lat_grid[j]/15)

    # Humidity decreasing with height
    q_field[i,j,k] = 0.015 * sigma_levels[k]
end

# Surface pressure field
ps_field = zeros(nx, ny)
for i in 1:nx, j in 1:ny
    ps_field[i,j] = 101325.0 - 2000.0 * sin(π * lat_grid[j]/15)
end

println("  ✓ Created temperature field (range: $(minimum(t_field)) to $(maximum(t_field)) K)")
println("  ✓ Created wind fields (u: $(minimum(u_field)) to $(maximum(u_field)) m/s)")
println("  ✓ Created humidity field (range: $(minimum(q_field)) to $(maximum(q_field)) kg/kg)")
println("  ✓ Created surface pressure (range: $(minimum(ps_field)) to $(maximum(ps_field)) Pa)")
println()

# ============================================================================
# Define observation locations
# ============================================================================

println("Step 2: Defining observation locations")
println("-" * "="^68)

nobs = 5

# Observation locations (lon, lat, pressure)
obs_locs = [
    -95.0  35.0  85000.0;  # Oklahoma City region, 850 hPa
    -90.0  40.0  70000.0;  # Illinois, 700 hPa
    -85.0  42.0  50000.0;  # Michigan, 500 hPa
    -92.0  38.0  92500.0;  # Missouri, surface
    -88.0  37.0  30000.0   # Upper air, 300 hPa
]

println("  Number of observations: $(nobs)")
for i in 1:nobs
    println("    Obs $i: Lon=$(obs_locs[i,1])°, Lat=$(obs_locs[i,2])°, P=$(obs_locs[i,3]/100) hPa")
end
println()

# ============================================================================
# Forward Operators: Compute model equivalents
# ============================================================================

println("Step 3: Computing forward operators")
println("-" * "="^68)

# Temperature observations
println("  Temperature observations:")
t_model = zeros(nobs)
for i in 1:nobs
    t_model[i] = temperature_forward(t_field, lon_grid, lat_grid, sigma_levels,
                                     obs_locs[i,1], obs_locs[i,2], obs_locs[i,3],
                                     ps_field)
    @printf("    Obs %d: %.2f K\n", i, t_model[i])
end
println()

# Wind observations (u component)
println("  Wind (u-component) observations:")
u_model = zeros(nobs)
for i in 1:nobs
    u_model[i] = wind_forward(u_field, v_field, lon_grid, lat_grid, sigma_levels,
                             obs_locs[i,1], obs_locs[i,2], obs_locs[i,3],
                             ps_field, :u, rotate_to_earth=false)
    @printf("    Obs %d: %.2f m/s\n", i, u_model[i])
end
println()

# Humidity observations
println("  Specific humidity observations:")
q_model = zeros(nobs)
for i in 1:nobs
    q_model[i] = humidity_forward(q_field, lon_grid, lat_grid, sigma_levels,
                                 obs_locs[i,1], obs_locs[i,2], obs_locs[i,3],
                                 ps_field, :specific)
    @printf("    Obs %d: %.6f kg/kg\n", i, q_model[i])
end
println()

# Surface pressure observations
println("  Surface pressure observations:")
ps_model = zeros(2)
for i in 1:2
    ps_model[i] = surface_pressure_forward(ps_field, lon_grid, lat_grid,
                                           obs_locs[i,1], obs_locs[i,2])
    @printf("    Obs %d: %.2f hPa\n", i, ps_model[i]/100)
end
println()

# ============================================================================
# Create synthetic observations and innovations
# ============================================================================

println("Step 4: Computing innovation vectors")
println("-" * "="^68)

# Add small perturbations to create "observations"
t_obs = t_model .+ randn(nobs) .* 0.5  # ±0.5 K observation error
u_obs = u_model .+ randn(nobs) .* 1.0  # ±1.0 m/s observation error
q_obs = q_model .+ randn(nobs) .* 0.001  # ±0.001 kg/kg observation error

# Innovations (observation - model)
t_innov = t_obs .- t_model
u_innov = u_obs .- u_model
q_innov = q_obs .- q_model

println("  Temperature innovations (O-B):")
for i in 1:nobs
    @printf("    Obs %d: %+.3f K\n", i, t_innov[i])
end
println()

println("  Wind innovations (O-B):")
for i in 1:nobs
    @printf("    Obs %d: %+.3f m/s\n", i, u_innov[i])
end
println()

# ============================================================================
# Adjoint Operators: Compute gradients
# ============================================================================

println("Step 5: Computing adjoint operators (gradients)")
println("-" * "="^68)

# Temperature gradient
t_gradient = zeros(nx, ny, nz)
ps_gradient_t = zeros(nx, ny)

for i in 1:nobs
    temperature_adjoint!(t_gradient, ps_gradient_t, t_innov[i],
                        lon_grid, lat_grid, sigma_levels,
                        obs_locs[i,1], obs_locs[i,2], obs_locs[i,3],
                        ps_field)
end

println("  Temperature gradient:")
@printf("    Norm: %.6e\n", norm(t_gradient))
@printf("    Max: %.6e\n", maximum(abs.(t_gradient)))
@printf("    Min: %.6e\n", minimum(abs.(t_gradient[t_gradient .!= 0])))
println()

# Wind gradient
u_gradient = zeros(nx, ny, nz)
v_gradient = zeros(nx, ny, nz)
ps_gradient_u = zeros(nx, ny)

for i in 1:nobs
    wind_adjoint!(u_gradient, v_gradient, ps_gradient_u, u_innov[i],
                 lon_grid, lat_grid, sigma_levels,
                 obs_locs[i,1], obs_locs[i,2], obs_locs[i,3],
                 ps_field, :u, rotate_to_earth=false)
end

println("  Wind (u-component) gradient:")
@printf("    Norm: %.6e\n", norm(u_gradient))
@printf("    Max: %.6e\n", maximum(abs.(u_gradient)))
println()

# Humidity gradient
q_gradient = zeros(nx, ny, nz)
ps_gradient_q = zeros(nx, ny)

for i in 1:nobs
    humidity_adjoint!(q_gradient, ps_gradient_q, q_innov[i],
                     lon_grid, lat_grid, sigma_levels,
                     obs_locs[i,1], obs_locs[i,2], obs_locs[i,3],
                     ps_field, :specific)
end

println("  Humidity gradient:")
@printf("    Norm: %.6e\n", norm(q_gradient))
@printf("    Max: %.6e\n", maximum(abs.(q_gradient)))
println()

# ============================================================================
# Verify Adjoint Property
# ============================================================================

println("Step 6: Verifying adjoint property")
println("-" * "="^68)

# For temperature operator
# ⟨H(x), y⟩ = ⟨x, Hᵀ(y)⟩

dot_forward = dot(t_model, t_innov)
dot_adjoint = dot(t_field[:], t_gradient[:])

println("  Temperature operator adjoint test:")
@printf("    ⟨H(x), y⟩  = %.12e\n", dot_forward)
@printf("    ⟨x, Hᵀ(y)⟩ = %.12e\n", dot_adjoint)
@printf("    Difference = %.6e\n", abs(dot_forward - dot_adjoint))

rel_error = abs(dot_forward - dot_adjoint) / max(abs(dot_forward), abs(dot_adjoint))
@printf("    Relative error = %.6e\n", rel_error)

if rel_error < 1e-10
    println("    ✓ ADJOINT TEST PASSED (machine precision)")
else
    println("    ⚠ Adjoint test: relative error = $(rel_error)")
end
println()

# ============================================================================
# Summary
# ============================================================================

println("="^70)
println("DEMO SUMMARY")
println("="^70)
println()
println("Successfully demonstrated:")
println("  ✓ Forward operators for temperature, wind, humidity, pressure")
println("  ✓ Interpolation from grid points to arbitrary observation locations")
println("  ✓ Innovation vector computation (O-B)")
println("  ✓ Adjoint operators for gradient computation")
println("  ✓ Adjoint consistency verification (⟨H(x),y⟩ = ⟨x,Hᵀ(y)⟩)")
println()
println("Total observations processed: $(nobs)")
println("Grid resolution: $(nx) x $(ny) x $(nz)")
println()
println("These operators are ready for use in variational data assimilation!")
println("="^70)
